from typing import ClassVar, no_type_check, override

import matplotlib.pyplot as plt
import numpy as np

from ...processors import ClusterStage
from ...utils import OutputManager
from ..interface import Marshaler

__all__ = ["ImageClusterMarshaler"]


class ImageClusterMarshaler(Marshaler[ClusterStage]):
    OUTPUT_NAME: ClassVar[str] = "kmeans_clusters.png"

    _output_manager: OutputManager
    _kmeans_config: dict
    _config: dict

    def __init__(self, output_manager: OutputManager, global_config: dict) -> None:
        self._output_manager = output_manager
        self._kmeans_config = global_config["clusterizer"]["kmeans"]
        self._config = global_config["clusterizer"]["kmeans"]["image"]

    @no_type_check
    @override
    def marshal(self, stage: ClusterStage) -> None:
        fig = plt.figure(figsize=(self._config["figure_width"], self._config["figure_height"]))
        ax = fig.add_subplot(111, projection="3d")

        # Draw the scatter plot for each cluster
        for i in range(self._kmeans_config["n_clusters"]):
            cluster_points = stage.points[stage.labels == i]
            # Add jitter effect to scatter points
            jittered_points = cluster_points.copy()
            if len(jittered_points) > 0:
                jittered_points += np.random.normal(0, self._kmeans_config["jitter_scale"], jittered_points.shape)
                # Ensure the jittered values are within the valid range
                jittered_points = np.clip(jittered_points, 0, 1)

            color = self._config["colors"][i % len(self._config["colors"])]
            ax.scatter(
                jittered_points[:, 0],
                jittered_points[:, 1],
                jittered_points[:, 2],
                s=self._config["marker_size"],
                c=color,
                alpha=self._config["alpha"],
                label=f"Cluster {i + 1}",
            )

        # Draw the cluster centers
        ax.scatter(
            stage.centers[:, 0],
            stage.centers[:, 1],
            stage.centers[:, 2],
            s=self._config["marker_size"] * self._config["centroid_size_multiplier"],
            c="black",
            marker=self._config["centroid_marker"],
            label="Centroids",
        )

        # Add labels and title
        ax.set_title("User Clusters (3D)")
        ax.set_xlabel("Neutral")
        ax.set_ylabel("Positive")
        ax.set_zlabel("Negative")

        # Adjust the view angle for better visibility of all axes
        ax.view_init(elev=20, azim=-35)

        # Ensure equal aspect ratio for better perception
        # Get the limits of the data
        x_limits = ax.get_xlim3d()
        y_limits = ax.get_ylim3d()
        z_limits = ax.get_zlim3d()

        # Get the range of each dimension
        x_range = abs(x_limits[1] - x_limits[0])
        y_range = abs(y_limits[1] - y_limits[0])
        z_range = abs(z_limits[1] - z_limits[0])

        # Get the middle of each dimension
        x_middle = np.mean(x_limits)
        y_middle = np.mean(y_limits)
        z_middle = np.mean(z_limits)

        # Set the same scale for all dimensions
        max_range = max([x_range, y_range, z_range])

        # Set the limits for each dimension
        ax.set_xlim3d([x_middle - max_range / 2, x_middle + max_range / 2])
        ax.set_ylim3d([y_middle - max_range / 2, y_middle + max_range / 2])
        ax.set_zlim3d([z_middle - max_range / 2, z_middle + max_range / 2])

        # Adjust padding to ensure axes labels are visible
        plt.tight_layout(pad=3.0)

        # Add legend with a better position
        ax.legend(loc="upper right", bbox_to_anchor=(1.15, 1))

        # Save the image
        plt.savefig(
            self._output_manager.filepath(ImageClusterMarshaler.OUTPUT_NAME),
            dpi=self._config["dpi"],
            bbox_inches="tight",
        )
        plt.close()
